Work around intermittent SM120 FP8 gradient corruption in RTC cast-transpose#3215
Conversation
Signed-off-by: AlbertYang514 <201034045+AlbertYang514@users.noreply.github.com>
Signed-off-by: AlbertYang514 <201034045+AlbertYang514@users.noreply.github.com>
github-actions
Bot
added
the
community-contribution
Greptile SummaryThis PR fixes a device-scope memory ordering race in the SM120 NVRTC-optimized
Confidence Score: 5/5Safe to merge; the kernel change is a single targeted fence with no effect on computed values or public APIs, and the test validates the specific failure mode. The production change is minimal — one __threadfence() call gated on an existing null-pointer check — and its placement (after all block-scope syncs, before the amax atomic) is correct. The validation and new regression test cover the specific failure scenario. The statically compiled fallback in transformer_engine/common/transpose/cast_transpose.cu (lines 203-209) has the same pre-fence pattern and would benefit from the same treatment if the RTC path is ever unavailable on SM120. Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant K as cast_transpose kernel (SM120)
participant GM as GPU Global Memory
participant Atom as atomicMaxFloat (amax_ptr)
participant DGEMM as dgrad GEMM kernel
K->>GM: store output_c (rowwise FP8)
K->>GM: store output_t (columnwise FP8, via shmem)
Note over K: __syncthreads() — block-scope barrier only
Note over K: __threadfence() — device-scope fence (NEW)
K->>Atom: reduce_max + atomicMaxFloat(amax_ptr, amax)
DGEMM->>GM: read FP8 outputs (guaranteed visible after fence)
DGEMM->>Atom: read amax_ptr (for next-step scaling)
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant K as cast_transpose kernel (SM120)
participant GM as GPU Global Memory
participant Atom as atomicMaxFloat (amax_ptr)
participant DGEMM as dgrad GEMM kernel
K->>GM: store output_c (rowwise FP8)
K->>GM: store output_t (columnwise FP8, via shmem)
Note over K: __syncthreads() — block-scope barrier only
Note over K: __threadfence() — device-scope fence (NEW)
K->>Atom: reduce_max + atomicMaxFloat(amax_ptr, amax)
DGEMM->>GM: read FP8 outputs (guaranteed visible after fence)
DGEMM->>Atom: read amax_ptr (for next-step scaling)
Reviews (2): Last reviewed commit: "test: relax SM120 dgrad norm bounds" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Signed-off-by: AlbertYang514 <201034045+AlbertYang514@users.noreply.github.com>
|
Hi @vthumbe1503 and @timmoon10, gentle ping when either of you has a chance. This is a narrowly scoped SM120 RTC cast-transpose ordering fix: one device-scope fence before the DelayedScaling amax atomic, plus a targeted regression test. It is currently mergeable, but has not received a human review yet. Could either of you please take a look, or point me to the current owner of this path? |
|
|
||
| // Reduce amax over block | ||
| if (amax_ptr != nullptr) { | ||
| // Order the global output stores before publishing the amax update. |
There was a problem hiding this comment.
I am actually unclear why fp8 output store needs to happen before amax compute. Isnt the dgrad compute stream ordered anyway? Why would dgrad recieve incorrect fp8 values and threadfence helps it?
There was a problem hiding this comment.
@vthumbe1503 — Your same-stream objection is correct, and the additional ablations do not support our original output-store/amax publication explanation. Seven fence placements were stable—including post-atomic, pre-return, and thread-0-only placements—while a system-scope atomic without the fence still failed, so I have rewritten the change as a timing/codegen-sensitive workaround. The current candidate retains the validated device-scope fence; the block-scope fence is documented only as experimental evidence. The exact root cause remains unknown, and the revised description now states that limitation explicitly.
AlbertYang514
changed the title
2 participants
Summary
This change retains a device-scope
__threadfence()in the optimized RTCcast_transposekernel as a validated workaround for intermittent FP8 gradient corruption observed on SM120. The evidence supports describing this as a timing/codegen-sensitive workaround; the exact root cause is still unknown.Observed issue
Under the affected checkpointed DelayedScaling workload, the forward loss remains normal, but backward can produce finite, extremely large gradients. The first affected layer and microbatch vary across runs. After the #3213 v2 checkpoint fix, the residual failure remains reproducible while FP8 ownership and update counts stay at the expected 1/1.
Change
The current code inserts
__threadfence()on the optimized RTCcast_transposepath. This device-scope fence is retained because it has been stable in the tested reproductions. It does not change the public API or intended mathematical operation.Investigation status
Nsight confirms that the optimized RTC cast-transpose kernel and the following dgrad GEMM execute on the same CUDA stream. Therefore, the original cross-kernel output-store/amax publication explanation is not supported by the execution model or the follow-up ablations. Seven tested fence placements were stable, including placements after the atomic operation, immediately before return, and under a thread-0-only condition. Replacing the fence with a system-scope atomic operation did not stabilize the no-fence variant, while a block-scope fence was also stable in the tested workload. Together, these results invalidate the earlier narrow publication-ordering explanation and instead point to a timing/codegen-sensitive effect. They do not identify the exact mechanism.
Validation
memcheck,racecheck,initcheck, andsynccheckcompleted without a reported error;initcheckwas also rerun without filtering. This is useful negative evidence, not proof that the kernel is free of races or memory-ordering defects.Relationship to #3213
The residual failure remains reproducible after #3213 v2 restores FP8 ownership/update counts to 1/1. This separates the observed residual phenomenon from the known duplicate-update pathology, but does not prove that every underlying mechanism is fully independent.
Scope and limitations
The validation is specific to the tested SM120 hardware, CUDA 13 toolchain, and checkpointed DelayedScaling workload. The block-scope fence is experimental evidence only; this PR retains the device-scope fence. The root cause remains unresolved, and no claim is made for other architectures or workloads. The change does not alter public APIs or intended math.